Skip to content

Add agent tool to manually regenerate a corpus icon#1938

Merged
JSv4 merged 4 commits into
mainfrom
claude/nice-thompson-mp7ND
Jun 6, 2026
Merged

Add agent tool to manually regenerate a corpus icon#1938
JSv4 merged 4 commits into
mainfrom
claude/nice-thompson-mp7ND

Conversation

@JSv4

@JSv4 JSv4 commented Jun 6, 2026

Copy link
Copy Markdown
Collaborator

Summary

Today the corpus icon/logo generator only runs once — automatically, at corpus-creation time, via the auto-branding flow (opencontractserver/corpuses/services/branding.py). This PR exposes that same generator as an on-demand agent tool so a user can ask the corpus chat agent to (re)generate a collection's icon at any point ("regenerate the icon", "give it a logo featuring a set of scales").

What's included

  • New tool regenerate_corpus_icon (opencontractserver/llms/tools/core_tools/corpus_branding.py)
    • Corpus-scoped, approval-gated, and creator-only — mirrors the existing corpus-row write surface (update_corpus_description / CorpusService.update_icon). The creator gate is checked up front (so we never spend an image-generation round-trip on a request that can't persist) and re-enforced by update_icon.
    • corpus_id / user_id are framework-injected (hidden from the LLM); the only LLM-visible argument is an optional additional_instructions styling hint, which is sanitised + length-capped before being folded into the image prompt.
  • New public helper aregenerate_corpus_logo (branding.py) shares the auto-branding primitives (_build_logo_prompt + agenerate_logo_image, with the deterministic PIL monogram fallback) and persists through the creator-gated CorpusService.update_icon. Unlike the create-time path it deliberately overwrites an existing icon and ignores auto_branding_enabled — a manual regeneration is an explicit request, not the best-effort create-time default. Corpus/user rows are re-loaded inside the save's database_sync_to_async boundary (mirroring _generate_logo).
  • _build_logo_prompt gains an optional additional_instructions argument (auto-branding path passes None → byte-identical behaviour).
  • Registry + agent wiring: registered in tool_registry.py (alias generate_corpus_icon; requires_approval + requires_corpus + requires_write_permission) and wired into the interactive corpus agent's authenticated toolset (pydantic_ai_agents.py) alongside update_corpus_description.
  • Constant CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS.

Permissioning

  • IDOR-safe: a corpus invisible to the user yields the same opaque "cannot access" message as a missing corpus.
  • Creator-only write (even a collaborator with guardian UPDATE cannot replace the icon), consistent with update_description/update_icon.
  • Only added to the corpus agent for authenticated users; the AI styling hint is sanitised like the title/description (a text-to-image model has no <user_content> fence).

Tests

opencontractserver/tests/test_corpus_icon_tool.py covers: the prompt builder's new hint handling (inclusion, blank-skip, quote neutralisation, length cap); the tool end-to-end (real monogram fallback persists a PNG, replaces an existing icon, the hint reaches the prompt); permission denials (anonymous, no-access, READ-only-but-not-creator, missing corpus); and registry wiring (flags, alias, approval gate firing through the wrapper).

Note: The backend test suite could not be executed in the authoring sandbox — building the Django test image requires PyPI, which is blocked by the environment's network policy (pip SSL cert failure). black and flake8 pass on all changed files, every file compiles, and the changelog fragment validates. CI will run the full suite here.

Changelog: changelog.d/regenerate-corpus-icon-tool.added.md.

Expose the corpus logo generator (previously only run once at corpus
creation by the auto-branding flow) as an on-demand agent tool so a user
can ask the corpus chat agent to (re)generate a collection's icon at any
time.

- New tool aregenerate_corpus_icon in
  opencontractserver/llms/tools/core_tools/corpus_branding.py: creator-only,
  approval-gated, corpus-scoped. Optional additional_instructions styling
  hint (sanitised + length-capped) is folded into the image prompt.
- New public helper aregenerate_corpus_logo in
  opencontractserver/corpuses/services/branding.py shares the auto-branding
  primitives (_build_logo_prompt + agenerate_logo_image) and persists via
  the creator-gated CorpusService.update_icon. Unlike the create-time path
  it deliberately overwrites an existing icon and ignores
  auto_branding_enabled.
- _build_logo_prompt gains an optional additional_instructions argument.
- Registered in tool_registry (alias generate_corpus_icon) and wired into
  the interactive corpus agent's authenticated toolset alongside
  update_corpus_description.
- Constant CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS.
- Tests in opencontractserver/tests/test_corpus_icon_tool.py.
@claude

claude Bot commented Jun 6, 2026

Copy link
Copy Markdown

Code Review — PR #1938: Add agent tool to manually regenerate a corpus icon

Overview

This PR adds a well-scoped regenerate_corpus_icon agent tool that lets a corpus creator ask the chat agent to (re)generate the corpus logo at any time. It cleanly reuses the existing auto-branding primitives, adds a sanitized styling hint, and is properly approval-gated. The architecture is sound; there are a few correctness and policy issues to address.


Issues

1. @pytest.mark.asyncio inside TransactionTestCase won't run correctly

opencontractserver/tests/test_corpus_icon_tool.py line 705

@pytest.mark.asyncio
async def test_requires_approval_fires(self):

Django's TransactionTestCase does not support async def test methods. pytest.mark.asyncio applied to a method on a TestCase subclass is silently ignored by pytest-django — the test will either be skipped or collected but never awaited correctly. The project pattern is to use async_to_sync for async work inside TestCase/TransactionTestCase:

def test_requires_approval_fires(self):
    registry = ToolFunctionRegistry.get()
    core_tool = registry.to_core_tool("regenerate_corpus_icon")
    assert core_tool is not None
    wrapper = PydanticAIToolWrapper(core_tool, inject_params={})
    callable_fn = wrapper.callable_function
    ctx = MagicMock()
    ctx.deps = PydanticAIDependencies(
        user_id=None, corpus_id=None, document_id=None, skip_approval_gate=False
    )
    ctx.tool_call_id = "test-call"
    with self.assertRaises(ToolConfirmationRequired) as cm:
        async_to_sync(callable_fn)(ctx, corpus_id=1, user_id=1)
    self.assertEqual(cm.exception.tool_name, "regenerate_corpus_icon")

2. additional_instructions_applied can return True when the hint was fully sanitized away

opencontractserver/llms/tools/core_tools/corpus_branding.py line 373

"additional_instructions_applied": bool(
    additional_instructions and additional_instructions.strip()
),

This flag is set based on the raw value before sanitize_plaintext_for_prompt runs. If the user passes a string that is non-empty/non-whitespace but whose entire content is stripped by the sanitizer (e.g., '"""' — pure quotes), additional_instructions_applied will be True while no hint actually appears in the prompt.

Either compute the flag after sanitization, or accept it as a best-effort "was a hint attempted" signal and document it as such. Given that the monogram path also ignores instructions, the former is more accurate.

3. Inline Tier-0 ORM call bypasses service layer (CLAUDE.md policy)

opencontractserver/llms/tools/core_tools/corpus_branding.py line 311

corpus = Corpus.objects.visible_to_user(user).filter(pk=corpus_id).first()

Per CLAUDE.md"Code with a user context … must reach models through opencontractserver/<app>/services/" and the note that LLM tools are policy-covered even though not yet mechanically enforced. CorpusService (or BaseService.filter_visible) should be used here rather than composing visible_to_user inline. update_corpus_description sets the precedent — check how that tool's authorizer resolves the corpus for the right model.

4. PR body contains prohibited Claude Code attribution

The PR description ends with:

_Generated by [Claude Code](https://claude.ai/code/session_015ZSJnpmusuPFVCkT5gNBj5)_

CLAUDE.md explicitly prohibits crediting Claude Code in PR messages ("Never credit Claude or Claude Code in commit messages, PR messages, comments, or any other artifacts"). Please remove this line before merging.


Minor Observations (non-blocking)

  • Patch target brittle if import is ever hoisted: _IMAGE_GEN_TARGET = "opencontractserver.utils.image_generation.agenerate_logo_image" is correct because the import is deferred inside aregenerate_corpus_logo. A comment next to _IMAGE_GEN_TARGET explaining this coupling would prevent a future refactor from silently producing always-passing mock tests.
  • No test for UPDATE-permission denial: test_non_creator_with_read_access_denied covers READ, but the PR description claims "even a collaborator with a guardian UPDATE grant cannot replace the icon." Adding a parallel test_non_creator_with_update_access_denied case would pin that guarantee explicitly.
  • user_id: int vs int | None: aregenerate_corpus_icon declares user_id: int but the test_anonymous_denied test passes user_id=None. The function body and authorizer handle None correctly; the type hint should be int | None to match reality (even if in practice the framework only injects ints for authenticated users).
  • @pytest.mark.django_db on TransactionTestCase subclass (line 685): redundant — TransactionTestCase already enables DB access. Harmless but noisy.

What's Done Well

  • Security: IDOR-safe via visible_to_user returning the same opaque error for both missing and inaccessible corpora; up-front creator gate prevents expensive image-generation round-trips on doomed requests; additional_instructions is sanitized + length-capped before prompt interpolation; the tool is wired only for authenticated users.
  • Architecture: Clean separation — the tool delegates to aregenerate_corpus_logo (branding service), which delegates to CorpusService.update_icon (authoritative creator gate). The DB re-fetch inside _save() handles mid-flight state changes gracefully.
  • Tests: Good coverage of the happy path, icon replacement, the hint-to-prompt wiring, and the full permission matrix (anonymous, no-access, read-only). The TransactionTestCase choice for the end-to-end tests is correct.
  • Changelog fragment: Detailed and properly placed.
  • Constant extraction: CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS in the constants file rather than a magic number.

- Import the concrete opencontractserver.users.models.User model class for
  the creator/other type annotations instead of the get_user_model() runtime
  variable, which mypy rejects as a type (valid-type) and pyupgrade unquotes
  under 'from __future__ import annotations'. Matches test_caml_review_tools.
- Guard FieldFile.name (str | None) and AsyncMock.await_args (Call | None)
  before attribute access; mark the intentional anonymous user_id=None call
  with type: ignore[arg-type].
@claude

claude Bot commented Jun 6, 2026

Copy link
Copy Markdown

Code Review — Add agent tool to manually regenerate a corpus icon

Overview

This PR exposes the existing corpus auto-branding generator as an on-demand agent tool (regenerate_corpus_icon). The implementation is well-structured: it reuses _build_logo_prompt/agenerate_logo_image through a new aregenerate_corpus_logo helper, is correctly approval-gated, creator-only, and IDOR-safe. The overall approach is sound.


Security

Prompt injection defence is solid. sanitize_plaintext_for_prompt + length-cap on additional_instructions before it's folded into the image prompt is exactly the right pattern, and the test test_instructions_quotes_are_neutralised verifies it.

IDOR-safe design is correct. Both "corpus not found" and "corpus not visible to user" hit the same visible_to_user(user).filter(pk=corpus_id).first() code-path and produce the same opaque PermissionError, making enumeration impossible.

Defence-in-depth on the creator gate (checked up front AND re-enforced by CorpusService.update_icon) is good practice.


Potential Bugs

1. @pytest.mark.asyncio on a TransactionTestCase method (test_corpus_icon_tool.py:711)

@pytest.mark.asyncio
async def test_requires_approval_fires(self):

TransactionTestCase is a Django TestCase subclass. Django handles async test methods natively (since 4.1) by running them inside asyncio.run(). Adding @pytest.mark.asyncio on top may cause double-wrapping or silent no-ops depending on the pytest-asyncio mode configured. The @pytest.mark.django_db on the class is also redundant — TransactionTestCase already handles DB access. Consider using a plain async def test_* without the asyncio marker, or extracting this test to a standalone @pytest.mark.asyncio function outside the class.

2. additional_instructions_applied can return True when sanitization strips the hint to empty (corpus_branding.py:373-375)

"additional_instructions_applied": bool(
    additional_instructions and additional_instructions.strip()
),

This checks whether the input was non-blank, not whether the sanitized hint actually made it into the prompt. A value consisting entirely of quotes (e.g. '"""') passes this check but is stripped to empty by sanitize_plaintext_for_prompt, so no style guidance is applied yet the response claims it was. The fix is to capture whether the sanitized hint is non-empty (either in _build_logo_prompt or by checking after sanitisation in the tool).


Code Quality / Style

3. Constant deferred-imported inside _build_logo_prompt (branding.py:175-177)

from opencontractserver.constants.corpus_branding import (
    CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS,
)

opencontractserver.constants.corpus_branding is a leaf constants module — no circular-import risk. Moving this to the module-level import block alongside the other top-level imports would improve readability, consistent with how CORPUS_LOGO_* constants are imported elsewhere.

4. PR footer violates project policy (PR description, last line)

_Generated by [Claude Code](https://claude.ai/code/...)_ — CLAUDE.md explicitly states: "Never credit Claude or Claude Code in commit messages, PR messages, comments, or any other artifacts." Please remove this line before merging.


Performance

The pre-generation permission check (before the up-to-90s image round-trip) is a well-considered optimisation. The database_sync_to_async boundaries are correctly placed, and the fresh DB reload inside _save correctly handles state changes that arrive during the slow image generation window.


Test Coverage

Coverage is good: prompt builder edge cases (blank, quote injection, length cap), happy path (real monogram fallback, icon replacement, hint propagation into prompt), and all denial branches (anonymous, no-access, read-only collaborator, missing corpus, approval gate). The mock target (opencontractserver.utils.image_generation.agenerate_logo_image) is correctly scoped given the deferred import pattern in aregenerate_corpus_logo.

One coverage note: test_creator_regenerates_icon_real_fallback relies on CORPUS_LOGO_GENERATION_ENABLED=False in test settings to avoid live API calls. There's no fallback assertion to verify this setting is in effect before the test runs, so a misconfigured environment would silently make real OpenAI requests. A guard like self.assertFalse(settings.CORPUS_LOGO_GENERATION_ENABLED) at the top of the test would make this assumption explicit and the failure obvious.


Summary

Two items worth addressing before merge:

  • test_requires_approval_fires — drop @pytest.mark.asyncio/@pytest.mark.django_db (Django handles async TransactionTestCase methods natively) or move the test to a standalone pytest function.
  • additional_instructions_applied — should reflect the sanitised result, not the raw input.

The core implementation (branding helper, tool, registry wiring, permission model) is clean and follows project conventions. Nice work on the prompt-injection hardening and defence-in-depth on the creator gate.

Place the concrete-User-model import in alphabetical order (between
types.enums and utils.*) as a bare import — the standalone explanatory
comment block sat out of order. Moved the rationale to a comment beside
the creator/other annotations instead.
@claude

claude Bot commented Jun 6, 2026

Copy link
Copy Markdown

Code Review: Add agent tool to manually regenerate a corpus icon

Overview

This PR exposes the existing corpus logo generator as an on-demand agent tool so a user can ask the corpus chat agent to regenerate the collection icon at any time. It is a well-scoped, single-purpose addition that reuses existing infrastructure cleanly.


Strengths

  • Double-gating is correct. The creator check runs up-front in _authorize_icon_regeneration (so no image-generation round-trip is wasted on a forbidden request) and is re-enforced by CorpusService.update_icon itself. Clear, intentional, and documented.
  • IDOR-safe errors. The "cannot access" branch returns the same opaque message whether the corpus is missing or merely invisible to the user, consistent with the project enumeration-prevention pattern.
  • Prompt injection hardening documented in-code. The security reasoning around sanitize_plaintext_for_prompt and the length cap is explained in the _build_logo_prompt docstring.
  • Deferred import in aregenerate_corpus_icon correctly avoids a circular import. The comment explains why.
  • Test coverage is comprehensive. Happy path, icon replacement, hint-reaches-prompt, all four denial cases (anonymous, no-access, read-but-not-creator, missing corpus), and registry/approval-gate wiring are all covered.
  • Fallback never raises. agenerate_logo_image always returns bytes (AI failure -> monogram), so aregenerate_corpus_logo will not throw unexpectedly on missing credentials.

Issues

1. CLAUDE.md violation: PR body credits Claude Code

Rule: "Never credit Claude or Claude Code in commit messages, PR messages, comments, or any other artifacts."

The PR description ends with _Generated by [Claude Code](...)_. Please remove that line.


2. @pytest.mark.django_db on a TransactionTestCase subclass is redundant

RegenerateCorpusIconRegistryTests(TransactionTestCase) already has full DB access through Django's own machinery. The @pytest.mark.django_db decorator is a no-op here and adds noise — it is intended for plain pytest functions/classes that do not inherit from Django's TestCase. Compare with AsyncImageToolsTestCase(TransactionTestCase) in test_image_tools.py, which does not carry the decorator. Suggest removing it.


3. Minor: constant imported inside _build_logo_prompt

CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS is a constant and cannot cause circular imports, so the deferred-import pattern is not necessary here. It was added alongside the already-deferred sanitize_plaintext_for_prompt import for consistency, but it would be cleaner at module level. Not blocking.


4. Minor: PIL imported inline in test_creator_regenerates_icon_real_fallback

from PIL import Image  # inside the test method body

PIL is a project-level dependency; it should be a top-level import consistent with every other import in the test file.


Summary

The core logic is sound: good security posture, correct async patterns, proper service-layer routing through CorpusService, and solid test coverage. The two actionable items before merge are:

  1. Remove the "Generated by Claude Code" attribution from the PR description (CLAUDE.md rule).
  2. Remove the redundant @pytest.mark.django_db decorator from RegenerateCorpusIconRegistryTests.

Items 3 and 4 are minor polish.

- branding.py: extract sanitize_logo_instruction_hint() and use it in
  _build_logo_prompt so the "hint applied?" answer derives from the sanitised
  value in one place.
- corpus_branding tool: report additional_instructions_applied from the
  sanitised hint (a value that collapses to empty — e.g. only quotes — no
  longer falsely reports as applied); annotate injected user_id as int | None
  to match the existing None guard and drop the now-stale test type: ignore.
- test_corpus_icon_tool.py: hoist the PIL import to module level, add an
  explicit settings.CORPUS_LOGO_GENERATION_ENABLED guard on the real-fallback
  test, and add a creator-only denial test for a collaborator holding a
  guardian UPDATE grant.

Kept the @pytest.mark.asyncio + class-level @pytest.mark.django_db markers:
they mirror test_image_tools.py, where the marker grants the async
TransactionTestCase method DB access under pytest-asyncio.
@claude

claude Bot commented Jun 6, 2026

Copy link
Copy Markdown

Code Review — PR #1938: Add agent tool to manually regenerate a corpus icon

Overall this is a well-structured, thoughtful PR. The intent is clear, the permissioning story is solid, and the test suite is thorough. A few items warrant attention before merge.


Issues

1. Inline visible_to_user in _authorize_icon_regeneration violates the service-layer policy

File: opencontractserver/llms/tools/core_tools/corpus_branding.py:328

corpus = Corpus.objects.visible_to_user(user).filter(pk=corpus_id).first()

Per CLAUDE.md: "Code with a user context … must reach models through opencontractserver/<app>/services/ — never compose visible_to_user … inline." The mechanical E001 check only scans config/graphql/, so this won't fail CI, but LLM tools are explicitly listed as user-context surfaces that the policy covers. The other corpus tool authorization helpers (_get_corpus_for_user, etc.) in sibling modules should already have this covered — worth checking if any of them can be reused here, or if CorpusService can expose a lookup method.

2. Double sanitization on the happy path

Files: corpus_branding.py:397 and branding.py:399

sanitize_logo_instruction_hint is called twice per tool invocation: once inside _build_logo_prompt (via aregenerate_corpus_logo) and again back in aregenerate_corpus_icon to compute additional_instructions_applied. The function is cheap, so this is not a performance problem, but it is logically redundant — the "applied" boolean could be computed from the return value of the prompt builder or by passing the sanitized hint through rather than re-sanitizing. A minor cleanliness nit.

3. Error message leaks user_id for the "user not found" path

File: opencontractserver/llms/tools/core_tools/corpus_branding.py:326

raise PermissionError(f"User {user_id} not found.")

This message surfaces the requesting user's PK. Because PermissionError propagates to the wrapper rather than being returned to the LLM, the string is never shown to the model — so there is no information-disclosure to the agent. However, it can appear in server logs, and the PR's own IDOR documentation emphasises returning opaque messages. Consider "Authentication error." or simply dropping the PK from this message for consistency.


Minor Observations

4. Creator comparison style divergence (cosmetic)

_authorize_icon_regeneration uses corpus.creator_id != user.id while CorpusService.update_icon uses corpus.creator_id != getattr(user, "id", None). Both are correct here (by the time user.id is accessed, user has been validated to be non-None), but the inconsistency is a future-reader puzzle. The getattr defensive form is the house style.

5. Lazy imports inside sanitize_logo_instruction_hint

File: branding.py:361-369

The two lazy imports inside this sync utility function (CORPUS_LOGO_ADDITIONAL_INSTRUCTIONS_MAX_CHARS, sanitize_plaintext_for_prompt) avoid a circular-import risk at module load — the comment chain documents the reason on aregenerate_corpus_logo. But sanitize_logo_instruction_hint itself is now a public export. It may be worth a one-line note in its docstring explaining why the imports are deferred (currently only the sibling function has that context).

6. Mid-flight deletion test coverage gap

The _save closure in aregenerate_corpus_logo handles the case where the corpus is deleted between image generation and save (returns a ServiceResult.failure), but there is no test exercising that path. This is admittedly tricky to test since it requires timing control, but a unit test with a mocked Corpus.objects.get raising DoesNotExist inside the async boundary would cover it.


What's working well

  • Approval-gate + creator-only write: The requires_approval=True flag plus the up-front creator check (failing before generating the image) and the re-enforcement inside update_icon are a defence-in-depth pattern that mirrors the existing corpus-row write surface cleanly.
  • IDOR-safe error framing: The "cannot access" branch returns the same opaque message for "doesn't exist" and "exists but you can't see it" — correct.
  • Prompt sanitization: The security note in _build_logo_prompt was explicitly extended to cover the new additional_instructions surface, and the test for quote-neutralisation (test_instructions_quotes_are_neutralised) gives confidence that sanitize_plaintext_for_prompt is actually doing the stripping.
  • Graceful operational failures: ValueError on mid-flight failure is correctly classified as an operational exception (the wrapper converts it to an error string for the LLM rather than crashing the turn).
  • Test coverage breadth: All four permission denial paths (anonymous, no-access, read-only non-creator, update-granted non-creator) are explicitly tested, as is registry wiring and the approval gate.
  • auto_branding_enabled is intentionally ignored: This is documented and correct — a manual regeneration is an explicit request, not a best-effort default.
  • Changelog fragment and constant in the right places: No magic numbers, fragment in changelog.d/.

@JSv4 JSv4 merged commit 8fbedae into main Jun 6, 2026
11 checks passed
@JSv4 JSv4 deleted the claude/nice-thompson-mp7ND branch June 6, 2026 18:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants